#include <linux/bitops.h>
#include <linux/minmax.h>
#include <linux/bitmap.h>

/*
 * Common helper for find_bit() function family
 * @FETCH: The expression that fetches and pre-processes each word of bitmap(s)
 * @MUNGE: The expression that post-processes a word containing found bit (may be empty)
 * @size: The bitmap size in bits
 */
#define FIND_FIRST_BIT(FETCH, MUNGE, size)                             \
    ({                                                                 \
        unsigned long idx, val, sz = (size);                           \
                                                                       \
        for (idx = 0; idx * BITS_PER_LONG < sz; idx++)                 \
        {                                                              \
            val = (FETCH);                                             \
            if (val)                                                   \
            {                                                          \
                sz = min(idx * BITS_PER_LONG + __ffs(MUNGE(val)), sz); \
                break;                                                 \
            }                                                          \
        }                                                              \
                                                                       \
        sz;                                                            \
    })

/*
 * Common helper for find_next_bit() function family
 * @FETCH: The expression that fetches and pre-processes each word of bitmap(s)
 * @MUNGE: The expression that post-processes a word containing found bit (may be empty)
 * @size: The bitmap size in bits
 * @start: The bitnumber to start searching at
 */
#define FIND_NEXT_BIT(FETCH, MUNGE, size, start)                      \
    ({                                                                \
        unsigned long mask, idx, tmp, sz = (size), __start = (start); \
                                                                      \
        if (unlikely(__start >= sz))                                  \
            goto out;                                                 \
                                                                      \
        mask = MUNGE(BITMAP_FIRST_WORD_MASK(__start));                \
        idx = __start / BITS_PER_LONG;                                \
                                                                      \
        for (tmp = (FETCH) & mask; !tmp; tmp = (FETCH))               \
        {                                                             \
            if ((idx + 1) * BITS_PER_LONG >= sz)                      \
                goto out;                                             \
            idx++;                                                    \
        }                                                             \
                                                                      \
        sz = min(idx * BITS_PER_LONG + __ffs(MUNGE(tmp)), sz);        \
    out:                                                              \
        sz;                                                           \
    })

unsigned long _find_next_and_bit(const unsigned long *addr1, const unsigned long *addr2,
                                 unsigned long nbits, unsigned long start)
{
    return FIND_NEXT_BIT(addr1[idx] & addr2[idx], /* nop */, nbits, start);
}

unsigned long find_next_bit(const unsigned long *addr, unsigned long size,
                            unsigned long start)
{
    unsigned long mask, idx, tmp, sz = size, __start = (start);

    if (__start >= sz)
        goto out;

    mask = BITMAP_FIRST_WORD_MASK(__start);
    idx = __start / BITS_PER_LONG;

    for (tmp = addr[idx] & mask; !tmp; tmp = addr[idx])
    {
        if ((idx + 1) * BITS_PER_LONG >= sz)
            goto out;
        idx++;
    }

    sz = min(idx * BITS_PER_LONG + __ffs(tmp), sz);

out:
    return sz;
}

/*
 * Find the first set bit in a memory region.
 */
unsigned long _find_first_bit(const unsigned long *addr, unsigned long size)
{
    return FIND_FIRST_BIT(addr[idx], /* nop */, size);
}

unsigned long find_first_bit(const unsigned long *addr, unsigned long size)
{
    if (small_const_nbits(size))
    {
        unsigned long val = *addr & GENMASK(size - 1, 0);

        return val ? __ffs(val) : size;
    }

    return _find_first_bit(addr, size);
}
